#include "ic3.h"
#include "base/aig.h"
#include "cktsat/cktcover.h"
#include "cktsat/satsolver.h"

#include <iostream>

//////////////////////////////
//  DEBUGGING INFO
//
void print_clause( const Clause & c , unsigned ns ){
   for( unsigned i = 0 ; i < ns ; ++ i )
   {
      char ch = '-';
      for( int j = 0 ; j < c.size() ; ++ j )
         if( i == (unsigned)var(c[j]) ){ ch = (sign(c[j])?'1':'0') ; break; }
      std::cout << ch;
   }
}
void print_states( const State s, unsigned ns ){
   for( uint64_t i=0 , j=0 , k=1 ; i < ns ; ((++i)&0x3F) ? (k+=k) : (k=1,j+=1) )
      std::cout << ((s[j]&k)?'1':'0');
}
static void debug_info( const State s , const Clause  & c , unsigned ns , unsigned i , unsigned k )
{
   std::cout <<"["<< i << ":" << k << "]\t";
   print_states(s,ns);
   std::cout <<"\t";
   print_clause(c,ns);
   std::cout << std::endl;
}

//////////////////////////////
//  Static Function
//
static unsigned b2w( unsigned b ){ return ( (b-1)>>6 ) +1 ;}

/**
@brief create a cover circuit with two frame, frame1: all input and state , output and po , frame 2 property used pi and po
*/
static void 
ckt_transform( StrashAIG & aig , CktCover & ckt )
{
   assert( aig._po.size() -1 == aig._andFlag - aig._ppiFlag );

   uint8_t * mrk = new uint8_t [aig.sz];
   AigLit  * lut = new AigLit [aig.sz]; 
   AigLit  * lut2= new AigLit [aig.sz];
   lut[0] = lut2[0] = 0;

   memset( mrk , 0 , sizeof( uint8_t ) * aig.sz );

   /// traverse second time frame from property to input
   ///
   mrk[ var( aig._po.last() ) ] = 3 ;
   for( unsigned i = aig.sz-1; i >= aig._andFlag ; -- i )
   {
      if( mrk[i] == 3 )
      {
         mrk[i] = 1;
         AigNode n = aig.data[i];
         if( mrk[var(n.fanin[0])] == 0 ) mrk[var(n.fanin[0])] = 3;
         if( mrk[var(n.fanin[1])] == 0 ) mrk[var(n.fanin[1])] = 3;
      } 
   }

   /// create Base , PI@0 , PPI@0 PI@1
   /// 
   {
      unsigned j = 1; 
      for( unsigned i = 1 ; i < aig._andFlag ; ++ i ) lut[i] = toAigLit( j++ , 0 );
      for( unsigned i = 1 ; i < aig._ppiFlag ; ++ i ) if( mrk[i] != 0 ) lut2[i] = toAigLit( j ++ , 0 );
      for( unsigned i = 1 ; i < j ; ++ i ) ckt.createBase();
      ckt._andFlag = j;
   }

   /// create AND create first time frame , create the ppi to second time frame , create second frame 
   ///
   for( unsigned i = aig._andFlag ; i < aig.sz ; ++ i )
   {
      AigNode n = aig[i];
      lut[i] = ckt.createAND( lut[var(n.fanin[0])]^sign(n.fanin[0]) , lut[var(n.fanin[1])]^sign(n.fanin[1]) );
   }
   for( unsigned i = aig._ppiFlag , j=0 ; i < aig._andFlag ; ++ i , ++ j )
      lut2[i] = lut[ var(aig._po[j]) ] ^ sign(aig._po[j]);
   for( unsigned i = aig._andFlag ; i < aig.sz ; ++ i )
   {
      if( mrk[i] == 0 ) continue;
      AigNode n = aig[i];
      lut2[i] = ckt.createAND( lut2[var(n.fanin[0])]^sign(n.fanin[0]) , lut2[var(n.fanin[1])]^sign(n.fanin[1]) );
   }

   /// create PO : PPO@0s , ~P@1, ~P@0 
   ///
   for( unsigned i = 0 ; i < aig._po.size()-1 ; ++i )
      ckt._po.push( lut[ var(aig._po[i]) ] ^ sign( aig._po[i] ) );
   ckt._po.push( lut2[ var(aig._po.last()) ] ^ sign( aig._po.last() ) );
   ckt._po.push( lut[ var(aig._po.last()) ] ^ sign( aig._po.last() ) );

   delete [] lut;
   delete [] lut2;
   delete [] mrk;
}


///////////////////////////////
//  LOCAL CLASS MEMBER
//
/**
@brief recording the clause at vec<Clause>
*/
class ClauseSaver : public SATSolver
{
   vec< Clause > & c;
   public:
   ClauseSaver( vec<Clause > & in ):c(in){}
   virtual ~ClauseSaver(){}
   virtual void addClause( vec<Lit> & x)
   {
      c.push();
      Clause & y = c.last();
      y.growTo( x.size() );
      for( unsigned i = 0 ; i < x.size() ; ++i ) //TODO using copy here
         y[i] = x[i];
   }
};
/**
@brief initialize the circuit to P and Finit and return sat(I&~P) , sat(I&T&~P')
*/
static unsigned init_assignLit( const StrashAIG & in , CktCover & ckt , Minisat::Lit * lits , Minisat::Lit * & rev )
{
   unsigned nvar =0 , j =0;

   Lit * l = ckt._lits;
   /// assign PPI literal 
   for( unsigned i = in._ppiFlag ; i < in._andFlag ; ++ i )
      lits[j++] = l[ i ] = Lit( nvar ++ , 0 );
   ckt.assignLit( nvar ); 

   nvar = ckt.nVars();
   rev = new Minisat::Lit [ nvar + ckt._po.size()];
   memset( rev , -1 , sizeof(Minisat::Lit) * (nvar + ckt._po.size()) );

   Minisat::Lit x;
   /// assign PPO literal 
   for( unsigned i = 0 ; i < ckt._po.size() -2 ; i ++ )
   {
#if 1
      x = l[ ::var(ckt._po[i]) ] ^ sign(ckt._po[i]);
      if( rev[ Minisat::var(x) ].x != -1 )
#endif
      x = Minisat::mkLit( nvar++ , 0 );
      lits[ j++ ] = x;
      rev[ Minisat::var(x) ] = Minisat::mkLit( i , Minisat::sign(x) );
   }

   int idx = ckt._po.size()-2;
   lits[j++] = l[ ::var(ckt._po[idx])] ^ sign(ckt._po[idx]) ^ 1;

   return nvar ;
}
static void init_assignCNF( CktCover & ckt , SATSolver & cs  , Minisat::Lit * lits, unsigned ns )
{
   Lit * l = ckt._lits;
   ckt.assignCNF( cs );
   {/// PPO assignment 
      for( unsigned i = 0 ; i < ckt._po.size() -2 ; i ++ )
      {
         Lit x = l[ ::var(ckt._po[i]) ] ^ sign(ckt._po[i]);
         Minisat::Lit m = lits[ i + ns ];
         Lit y = Lit( var(m),sign(m));
         if( x == y ) continue;
         cs.addClause( x,~y );
         cs.addClause( ~x,y );
      }
   }
}
static void state_dependency( const CktCover & ckt , SATSolver & cs )
{
/// construct registr relationship
   const vec<AigLit> & po = ckt._po;
   std::tr1::unordered_map< unsigned , unsigned > polit;
   std::tr1::unordered_map< unsigned , unsigned >::iterator iter;
   for( unsigned i = 0 ; i < po.size()-2 ; ++ i )
   {
      iter = polit.find( var(po[i]) );
      if( iter == polit.end() )
         polit.insert( std::make_pair( var(po[i]) , i ) );
      else
      {
         Lit a = Lit( iter->second , 0 );
         Lit b = Lit( i , sign( po[iter->second] ) ^ sign( po[i] )  );
         if( sign( po[iter->second] ) == sign( po[i] ) ) // could not insert 0 , 0
            cs.addClause( a , ~b );
         cs.addClause( ~a,  b ); 
      }
   }
}
bool IC3::init( StrashAIG & in )
{
   _nstates = in._andFlag - in._ppiFlag;
   StateAllocator::init_stalloc( b2w(_nstates) );
   _lits = new Minisat::Lit[ _nstates+_nstates+1 ];

   CktCover ckt;
   ckt_transform( in , ckt );
   ckt.init(false); // no phase simplify 
   ckt.covering();
  
   /// assign Literal , from state and record the literal to _lits
   Lit * l = ckt._lits;
   _nvar = init_assignLit( in , ckt , _lits ,_rlits );
   
   /// dump CNF 
   Finit.resize( ckt.nClause() + 1 + _nstates*2 ); Finit.clear();
   ClauseSaver cs(Finit);
   init_assignCNF( ckt , cs ,  _lits , _nstates );


   MinisatIC3 * f0 = addFrame();
   {/// first fram initial value 
      for( unsigned i = 0 ; i < _nstates ; ++ i )
      { 
         Clause * c = new Clause(1);
         (*c)[0] = ~ _lits[i];
         addClause( 0 , 0 , c );
      }
   }

   /// solve the initial resoult
   Minisat::Lit p = l[ ::var(ckt._po.last() ) ] ^ sign( ckt._po.last() ) ^ 1 ;
   if( f0->solve( ~p ) ){ /*cout << "0 frame sat" <<endl;*/ return true;}

   /// insert first frame property 
   ((SATSolver & )cs).addClause( Lit( Minisat::var(p),sign(p)) );
   f0->addClause( p ); 

   /// solver second frame 
   if( f0->solve( ~ _lits[ _nstates + _nstates ] ) ) { /*cout << "1 frame sat" <<endl;*/ return true;}

   /// insert state dependency without blocking initial value 
   state_dependency( ckt , cs );
   
   std::cout << "initial #sts " << _nstates << " #cls " << Finit.size() << " #var " << _nvar << std::endl;
   
   addFrame();
   return false;
}
IC3::~IC3()
{
   delete [] _lits;
   delete [] _rlits;

   for( unsigned i = 0 ; i < F.size() ; ++ i )
      delete F[i];
}
bool IC3::prove()
{
   for( int k = 1 ; k < INT_MAX ; ++ k )
   {
      if( ! strengthen( k ) )
         return false;
      
      propagateClause( k );
      remove_redundant_clause();
      for( int i = 1 ; i <= k ; ++ i )
         if( _clauses[i].size() == 0 )
            return true;
   }
   return false;
}
bool IC3::strengthen( int k )
{
   ic3info();

   addFrame();
   while( F[k]->sat() )
   {
      State s = stalloc();
      F[k]->witness( s );

      int n = inductivelyGeneralize( s , k-2 , k );
      States set( b2w(_nstates) );
      set.add( s , n+1 );
      if( pushGeneralization( set , k ) == false )
         return false;
      reset_stalloc();
   }
   return true;
}
void IC3::propagateClause( int k )
{
   for( int i = 0 ; i <= k ; ++ i )
   {
      unsigned sz = _clauses[i].size();
      for( unsigned j = 1 ; j < sz ; ++ j )
      {
         const Clause * c = clause( i , j );
         if( c == NULL ) continue;
         if( ! F[i]->sat( *c ) )
         {
            F[i+1]->addClause( *c );
            shift_clause( i , j );   
         }
      }
   }
}
int IC3::inductivelyGeneralize( State s , int min , int k )
{
   //assert( min < 0 || isInductiveInvariant( min , s ) );
   if( (min < 0) && F[0]->sat(s ) )
      return -1;

   for( int i = std::max( 1 , min+1) ; i <= k ; ++ i )
      if( F[i]->sat(s) )
      {
         //if( i-1 == min ) F[i-1]->sat(s); // querry a unsat 
         if( i-1 != min )
            generateClause( s , i-1 , k );
         return i-1;
      }
   generateClause( s , k , k );
   return k;
}
void IC3::generateClause( State s , int i , int k )
{
   Clause * c = new Clause();
   F[i]->subclause_inductive_relative(s,*c);
   debug_info( s , *c , _nstates , i , k );
   if( addClause( 0 , i+1 , c ) == false)delete c;
}
bool IC3::pushGeneralization( States & sset , int k )
{
   while( true )
   {
      sset.extract();
      int n   = sset.level();
      if( n == 0 ) return false;
      if( n >  k ) return true;

      State s = sset.state();;
      if( F[n]->sat(s) )
      {
         State p = stalloc();
         F[n]->witness( p );
         int m = inductivelyGeneralize( p , n-2 , k );
         if( sset.add( p , m+1 ) == false ) resume_stalloc();

      }else
      {
         int m = inductivelyGeneralize( s , n , k );
         sset.modify( m+1 );
      }
   }
}
/// TODO ClauseKeaper protected remove 
bool IC3::addClause( unsigned s , unsigned e , const Clause * c )
{
   CKHash::iterator iter = _hash.find( c );
   if( iter != _hash.end() )
   {
      clause_t & ct = _keaper[ iter->second ];
      
      if( ct._lvl == e )
      {
//         assert( false );
         return false ;
      }

      for( unsigned i = ct._lvl+1 ; i <= e ; ++ i)
         F[i]->addClause( *c );

      ct._lvl = e;
      _clauses[e].push( iter->second );
      return false;
   }else
   {
      add_clause( c , e ); // add to keaper 
      for( unsigned i = s ; i <= e ; ++ i )
         F[i]->addClause( *c );
      return true;
   }
}
void IC3::ic3info()
{
   std::cout << F.sz-1 << "\t";
   std::cout << "clauses:" << _keaper.size()<< " ";
   for( unsigned i = 0 ; i < _clauses.sz ; ++ i )
      std::cout << _clauses[i].sz<< " ";
   std::cout<<std::endl;
}
